import random

import utils.detection_3d.solder_splash as solder_splash
import numpy as np
import cv2
import os
import json


def add_tin_in_area(input_npy, input_fly, yl, yr, xl, xr, mask):
    ys, xs = mask.shape
    fly_y, fly_x = input_fly.shape
    for i in range(100):
        yrandom = random.randint(0, ys - 1)
        xrandom = random.randint(0, xs - 1)
        if yrandom + fly_y <= ys and xrandom + fly_x <= xs and np.sum(
                mask[yrandom:yrandom + fly_y, xrandom:xrandom + fly_x]) == fly_y * fly_x:
            print("area find success, (", yrandom, ", ", xrandom, ")")
            break
        else:
            if i < 99:
                print("--- generate try again, round ", i + 1)
            else:
                print("--- area generate failed")
                return input_npy
    fly_nan_index = np.isnan(input_fly)
    used_area = input_npy[yl + yrandom:yl + yrandom + fly_y, xl + xrandom:xl + xrandom + fly_x]
    input_fly[fly_nan_index] = 0
    used_area[~fly_nan_index] = 0
    input_fly[~fly_nan_index] = input_fly[~fly_nan_index] + max(0, np.mean(used_area[fly_nan_index]))
    add_area = input_fly + used_area
    # solder_splash.show_3d(input_fly)
    # solder_splash.show_3d(add_area)
    input_npy[yl + yrandom:yl + yrandom + fly_y, xl + xrandom:xl + xrandom + fly_x] = add_area
    return input_npy


def calculator_tin_diff_with_mask(x_points, y_points, z_points, std_npy, npy, start_point, mask, area, tin_dir,
                                  pic=None, target_folder=None, filename=None):
    """
    利用mask过滤
    Args:
        x_points:
        y_points:
        z_points:
        std_npy:
        npy:
        start_point:

    Returns:

    """
    # origin_nan = np.isnan(npy)
    # 计算底面
    bottom_down, bottom_up, bottom_left = solder_splash.get_bottom_value(x_points, y_points, z_points)
    Z_BOTTOM = (bottom_down + bottom_up + bottom_left) / 3  # PCD面的大致底面高度
    x_points, y_points, z_points = solder_splash.data_filter_by_area(x_points, y_points, z_points,
                                                                     x_min=start_point[1] / solder_splash.RADIO,
                                                                     x_max=(start_point[1] + std_npy.shape[
                                                                         1]) / solder_splash.RADIO,
                                                                     y_min=start_point[0] / solder_splash.RADIO,
                                                                     y_max=(start_point[0] + std_npy.shape[
                                                                         0]) / solder_splash.RADIO)
    real_bottom = solder_splash.get_real_bottom(Z_BOTTOM, z_points)

    # 滑窗计算结果
    result, y, x = solder_splash.move_window(std_npy, npy - real_bottom, start_point)
    print("result: ", str(result), " (", str(x), ", ", str(y), ")")
    tin_data = np.load(tin_dir)
    new_npy = add_tin_in_area(npy - real_bottom, tin_data, y, y + std_npy.shape[0], x, x + std_npy.shape[1], mask)
    return new_npy


def get_tin_file(files):
    ind = random.randint(0, len(files) - 1)
    for i in range(ind, len(files)):
        if files[i][-3:] == "npy":
            return files[i]
    for i in range(ind, -1, -1):
        if files[i][-3:] == "npy":
            return files[i]


def add_fly_tin(input_npy, target_folder, tin_folder, add_area, addition_name=None):
    global_cfg = "./global_cfg.json"
    with open(global_cfg, 'r') as f:
        config = json.load(f)
        AREA_43 = config['std-43-area']
        AREA_44 = config['std-44-area']
        AREA_45 = config['std-45-area']
        area_threshold = config['fly-tin-threshold']
    tin_file = os.listdir(tin_folder)
    print(input_npy)
    # 清理边缘
    data = solder_splash.clean_height(input_npy)
    # 求pcd板底面深度
    xyz_points = solder_splash.parse_npy_(data)  # 得到x，y，z坐标
    x_points, y_points, z_points = xyz_points

    # 读取标准件区域
    std_4_3 = np.load(solder_splash.STD_4_3)
    std_4_4 = np.load(solder_splash.STD_4_4)
    std_4_5 = np.load(solder_splash.STD_4_5)

    # 亮度图引入，意图清理无效nan点
    pic = solder_splash.clean_light(input_npy[:-4] + "_meta.png")
    pic = cv2.resize(pic, (data.shape[1], data.shape[0]))
    pic = cv2.cvtColor(pic, cv2.COLOR_BGR2GRAY)

    # solder_splash.show_3d(data)

    if add_area[0]:
        mask43 = np.load("./utils/detection_3d/mask_4-3.npy")
        ind = get_tin_file(tin_file)
        print("used fly tin: ", ind)
        filename = input_npy[:-4] + "_add_tin43.npy"
        data = calculator_tin_diff_with_mask(x_points, y_points, z_points, std_4_3, data, (AREA_43[2], AREA_43[0]),
                                             mask43, "43",
                                             tin_dir=os.path.join(tin_folder, ind), pic=pic,
                                             target_folder=target_folder, filename=filename)
        # show_3d(std_4_3, "std_4-3")

    # solder_splash.show_3d(data)

    if add_area[1]:
        filename = input_npy[:-4] + "_add_tin44.npy"
        ind = get_tin_file(tin_file)
        print("used fly tin: ", ind)
        mask44 = np.load("./utils/detection_3d/mask_4-4.npy")
        data = calculator_tin_diff_with_mask(x_points, y_points, z_points, std_4_4, data, (AREA_44[2], AREA_44[0]),
                                             mask44, "44",
                                             tin_dir=os.path.join(tin_folder, ind), pic=pic,
                                             target_folder=target_folder, filename=filename)
        # show_3d(std_4_4, "std_4-4")

    # solder_splash.show_3d(data)

    if add_area[2]:
        filename = input_npy[:-4] + "_add_tin45.npy"
        ind = get_tin_file(tin_file)
        print("used fly tin: ", ind)
        mask45 = np.load("./utils/detection_3d/mask_4-5.npy")
        data = calculator_tin_diff_with_mask(x_points, y_points, z_points, std_4_5, data, (AREA_45[2], AREA_45[0]),
                                             mask45, "45",
                                             tin_dir=os.path.join(tin_folder, ind), pic=pic,
                                             target_folder=target_folder, filename=filename)
    if addition_name is None:
        result_dir = os.path.join(target_folder, os.path.basename(input_npy)[:-4] + "_add_tin.npy")
    else:
        result_dir = os.path.join(target_folder,
                                  os.path.basename(input_npy)[:-4] + "_add_tin_" + addition_name + ".npy")
    np.save(result_dir, data)
    print("add fly tin file: ", result_dir)
    # solder_splash.show_3d(data)
    data = solder_splash.clean_height(result_dir)
    # 求pcd板底面深度
    xyz_points = solder_splash.parse_npy_(data)  # 得到x，y，z坐标
    x_points, y_points, z_points = xyz_points
    r43, y43, x43 = solder_splash.calculator_tin_diff_with_mask(x_points, y_points, z_points, std_4_3, data,
                                                                (AREA_43[2], AREA_43[0]), mask43, area="43",
                                                                target_folder=target_folder,
                                                                filename=os.path.basename(result_dir))
    # show_3d(std_4_3, "std_4-3")
    r44, y44, x44 = solder_splash.calculator_tin_diff_with_mask(x_points, y_points, z_points, std_4_4, data,
                                                                (AREA_44[2], AREA_44[0]), mask44, area="44",
                                                                target_folder=target_folder,
                                                                filename=os.path.basename(result_dir))
    # show_3d(std_4_4, "std_4-4")
    r45, y45, x45 = solder_splash.calculator_tin_diff_with_mask(x_points, y_points, z_points, std_4_5, data,
                                                                (AREA_45[2], AREA_45[0]), mask45, area="45",
                                                                target_folder=target_folder,
                                                                filename=os.path.basename(result_dir))